from transformers import BertTokenizer, BertForSequenceClassification
import torch
import torch.nn.functional as F
import pandas as pd

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. 加载训练好的模型和分词器
model_dir = r"E:\AI\models\bert\checkpoint\bert-base-chinese"  # 你的模型存放目录
tokenizer = BertTokenizer.from_pretrained(model_dir)
model = BertForSequenceClassification.from_pretrained(model_dir)
model.to(device)

# 2. 加载类别标签
label_file = f"data/label.txt"
with open(label_file, 'r', encoding='utf-8') as f:
    labels = [line.strip() for line in f]


# 3. 定义分类函数
def classify_text(text):
    # 对输入文本进行分词
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256).to(device)

    # 使用模型进行推理
    model.eval()  # 设定模型为评估模式
    with torch.no_grad():  # 不需要计算梯度
        outputs = model(**inputs)

    # 获取logits
    logits = outputs.logits
    # print("Logits shape:", logits.shape)
    # print("Logits:", logits)

    # 计算每个类别的概率
    probabilities = F.softmax(logits, dim=-1)
    # print("Probabilities:", probabilities)

    top_values, top_indices = torch.topk(probabilities, k=3, dim=-1)

    print("Top 3 values:", top_values)
    print("Top 3 indices:", top_indices)

    # 获取分数最高的类别及其概率
    predicted_class_id = torch.argmax(probabilities, dim=-1).item()
    predicted_probability = probabilities[0, predicted_class_id].item()

    # print("Predicted class ID:", predicted_class_id)
    # print("Predicted class probability:", predicted_probability)

    # 获取logits并找到分数最高的类别
    # print(outputs.logits.shape)
    # logits = outputs.logits
    # print(logits)
    # predicted_class_id = torch.argmax(logits, dim=-1).item()
    # print(predicted_class_id)


    # 返回预测的类别
    return labels[predicted_class_id], predicted_probability


def eval_all(file_path):
    df = pd.read_csv(file_path, sep='\t', header=None)
    acc_number = 0
    for i in range(df.shape[0]):
        text = df.iloc[i, 0]
        label = df.iloc[i, 1]
        predicted_label, predicted_probability = classify_text(text)
        print(f'{i}:  right: {label==predicted_label}, predicted_probability: {predicted_probability}, text: {text}, label: {label}, predicted_label: {predicted_label}')
        if predicted_label == label:
            acc_number += 1
    print(f'total: {df.shape[0]}, acc:{acc_number}, accuracy: {acc_number / df.shape[0]}')

# 4. 测试推理
# test_text = "后保险杠加强件缺失一汽丰田召回超13万辆皇冠汽车"
# predicted_label, probability = classify_text(test_text)
# print(f"预测类别: {predicted_label}")
eval_all('data/train.txt')



if __name__ == '__main__':
    pass